from typing import Tuple, Optional

import cv2
import torch
from torch import Tensor, LongTensor

from models.CDO.resnet import *
import numpy as np
import os
from models.CDO.hrnet.hrnet import HRNet_
from torch.nn import functional as F
import torch.nn as nn
import loguru
from scipy.ndimage import gaussian_filter
import math
import logging
import torchvision.transforms.functional as TF

from torch.nn.parameter import Parameter
valid_resnet_backbones = ['resnet18', 'resnet34', 'resnet50', 'wide_resnet50_2']
valid_hrnet_backbones = ['hrnet18', 'hrnet32', 'hrnet48']

from efficientnet_pytorch.model import MemoryEfficientSwish

# 定义一个通过卷积和激活函数生成注意力图的模块
class AttnMap(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.act_block = nn.Sequential(
            nn.Conv2d(dim, dim, 1, 1, 0),  # 1x1卷积用于调整通道数
            MemoryEfficientSwish(),  # 使用MemoryEfficientSwish作为激活函数
            nn.Conv2d(dim, dim, 1, 1, 0)   # 再次1x1卷积
        )

    def forward(self, x):
        return self.act_block(x)

# 定义高效注意力机制的主体模块
class EfficientAttention(nn.Module):
    def __init__(self, dim, num_heads=8, group_split=[4, 4], kernel_sizes=[5], window_size=4,
                 attn_drop=0., proj_drop=0., qkv_bias=True):
        super().__init__()
        # 参数初始化和定义
        assert sum(group_split) == num_heads  # 确保分组数量之和等于头的数量
        assert len(kernel_sizes) + 1 == len(group_split)  # 确保核大小列表加一等于分组数量
        self.dim = dim  # 输入通道数
        self.num_heads = num_heads  # 注意力头的数量
        self.dim_head = dim // num_heads  # 每个头的维度
        self.scalor = self.dim_head ** -0.5  # 缩放因子
        self.kernel_sizes = kernel_sizes  # 核大小列表
        self.window_size = window_size  # 窗口大小
        self.group_split = group_split  # 分组列表
        # 根据核大小和分组定义卷积层、注意力映射层和QKV层
        convs = []
        act_blocks = []
        qkvs = []
        for i in range(len(kernel_sizes)):
            kernel_size = kernel_sizes[i]
            group_head = group_split[i]
            if group_head == 0:
                continue
            convs.append(nn.Conv2d(3 * self.dim_head * group_head, 3 * self.dim_head * group_head, kernel_size,
                                   1, kernel_size // 2, groups=3 * self.dim_head * group_head))
            act_blocks.append(AttnMap(self.dim_head * group_head))
            qkvs.append(nn.Conv2d(dim, 3 * group_head * self.dim_head, 1, 1, 0, bias=qkv_bias))
        if group_split[-1] != 0:
            # 对最后一个全局注意力头的定义
            self.global_q = nn.Conv2d(dim, group_split[-1] * self.dim_head, 1, 1, 0, bias=qkv_bias)
            self.global_kv = nn.Conv2d(dim, group_split[-1] * self.dim_head * 2, 1, 1, 0, bias=qkv_bias)
            self.avgpool = nn.AvgPool2d(window_size, window_size) if window_size != 1 else nn.Identity()

        # 将定义的模块注册为子模块
        self.convs = nn.ModuleList(convs)
        self.act_blocks = nn.ModuleList(act_blocks)
        self.qkvs = nn.ModuleList(qkvs)
        self.proj = nn.Conv2d(dim, dim, 1, 1, 0, bias=qkv_bias)  # 输出投影层
        self.attn_drop = nn.Dropout(attn_drop)  # 注意力dropout
        self.proj_drop = nn.Dropout(proj_drop)  # 投影dropout

    # 高频注意力处理函数
    def high_fre_attntion(self, x: torch.Tensor, to_qkv: nn.Module, mixer: nn.Module, attn_block: nn.Module):
        '''
        x: (b c h w)
        '''
        b, c, h, w = x.size()
        qkv = to_qkv(x)  # (b (3 m d) h w)
        qkv = mixer(qkv).reshape(b, 3, -1, h, w).transpose(0, 1).contiguous()  # (3 b (m d) h w)
        q, k, v = qkv  # (b (m d) h w)
        attn = attn_block(q.mul(k)).mul(self.scalor)
        attn = self.attn_drop(torch.tanh(attn))
        res = attn.mul(v)  # (b (m d) h w)
        return res

    # 低频注意力处理函数
    def low_fre_attention(self, x: torch.Tensor, to_q: nn.Module, to_kv: nn.Module, avgpool: nn.Module):
        '''
        x: (b c h w)
        '''
        b, c, h, w = x.size()

        q = to_q(x).reshape(b, -1, self.dim_head, h * w).transpose(-1, -2).contiguous()  # (b m (h w) d)
        kv = avgpool(x)  # (b c h w)
        kv = to_kv(kv).view(b, 2, -1, self.dim_head, (h * w) // (self.window_size ** 2)).permute(1, 0, 2, 4,
                                                                                                 3).contiguous()  # (2 b m (H W) d)
        k, v = kv  # (b m (H W) d)
        attn = self.scalor * q @ k.transpose(-1, -2)  # (b m (h w) (H W))
        attn = self.attn_drop(attn.softmax(dim=-1))
        res = attn @ v  # (b m (h w) d)
        res = res.transpose(2, 3).reshape(b, -1, h, w).contiguous()
        return res

    # 模块的前向传播
    def forward(self, x: torch.Tensor):
        '''
        x: (b c h w)
        '''
        res = []
        for i in range(len(self.kernel_sizes)):
            if self.group_split[i] == 0:
                continue
            res.append(self.high_fre_attntion(x, self.qkvs[i], self.convs[i], self.act_blocks[i]))
        if self.group_split[-1] != 0:
            res.append(self.low_fre_attention(x, self.global_q, self.global_kv, self.avgpool))
        return self.proj_drop(self.proj(torch.cat(res, dim=1)))
def conv_with_padding(in_planes, out_planes, kernelsize, stride=1, dilation=1, bias=False, padding = None):
    if padding is None:
        padding = kernelsize//2
    return nn.Conv2d(in_planes, out_planes, kernel_size=kernelsize, stride=stride, dilation=dilation, padding=padding, bias=bias)


def conv_init(conv, act='linear'):
    r"""
    Reproduces conv initialization from DnCNN
    """
    n = conv.kernel_size[0] * conv.kernel_size[1] * conv.out_channels
    conv.weight.data.normal_(0, math.sqrt(2. / n))

def batchnorm_init(m, kernelsize=3):
    r"""
    Reproduces batchnorm initialization from DnCNN
    """
    n = kernelsize**2 * m.num_features
    m.weight.data.normal_(0, math.sqrt(2. / (n)))
    m.bias.data.zero_()

def make_activation(act):
    if act is None:
        return None
    elif act == 'relu':
        return nn.ReLU(inplace=True)
    elif act == 'tanh':
        return nn.Tanh()
    elif act == 'leaky_relu':
        return nn.LeakyReLU(inplace=True)
    elif act == 'softmax':
        return nn.Softmax()
    elif act == 'linear':
        return None
    else:
        assert(False)
def make_net(nplanes_in, kernels, features, bns, acts, dilats, bn_momentum = 0.1, padding=None):
    r"""
    :param nplanes_in: number of of input feature channels
    :param kernels: list of kernel size for convolution layers
    :param features: list of hidden layer feature channels
    :param bns: list of whether to add batchnorm layers
    :param acts: list of activations
    :param dilats: list of dilation factors
    :param bn_momentum: momentum of batchnorm
    :param padding: integer for padding (None for same padding)
    """

    depth = len(features)
    assert(len(features)==len(kernels))

    layers = list()
    for i in range(0,depth):
        if i==0:
            in_feats = nplanes_in
        else:
            in_feats = features[i-1]

        elem = conv_with_padding(in_feats, features[i], kernelsize=kernels[i], dilation=dilats[i], padding=padding, bias=not(bns[i]))
        conv_init(elem, act=acts[i])
        layers.append(elem)

        if bns[i]:
            elem = nn.BatchNorm2d(features[i], momentum = bn_momentum)
            batchnorm_init(elem, kernelsize=kernels[i])
            layers.append(elem)

        elem = make_activation(acts[i])
        if elem is not None:
            layers.append(elem)

    return nn.Sequential(*layers)

class noisePrint_extractor(nn.Module):
    def __init__(self):
        super().__init__()
        num_levels = 17
        out_channel = 1
        self.noiseprint = make_net(3, kernels=[3, ] * num_levels,
                                   features=[64, ] * (num_levels - 1) + [out_channel],
                                   bns=[False, ] + [True, ] * (num_levels - 2) + [False, ],
                                   acts=['relu', ] * (num_levels - 1) + ['linear', ],
                                   dilats=[1, ] * num_levels,
                                   bn_momentum=0.1, padding=1)

        noiseprint_path = "/home/wc/disk1/MMFusion/pretrained/noiseprint/np++.pth"
        np_weights = noiseprint_path
        assert os.path.isfile(np_weights)
        dat = torch.load(np_weights, map_location=torch.device('cpu'))
        logging.info(f'Noiseprint++ weights: {np_weights}')
        self.noiseprint.load_state_dict(dat)

        self.noiseprint.eval()
        for param in self.noiseprint.parameters():
            param.requires_grad = False

    def forward(self,x):
        out = self.noiseprint(x)
        if out.size()[-3] == 1:
            out = torch.tile(out, (3, 1, 1))
        return out

def _grid2seq(x: Tensor, region_size: Tuple[int], num_heads: int):
    B, C, H, W = x.size()
    region_h, region_w = H // region_size[0], W // region_size[1]
    x = x.view(B, num_heads, C // num_heads, region_h, region_size[0], region_w, region_size[1])
    x = torch.einsum('bmdhpwq->bmhwpqd', x).flatten(2, 3).flatten(-3, -2)  # (bs, nhead, nregion, reg_size, head_dim)
    return x, region_h, region_w


def _seq2grid(x: Tensor, region_h: int, region_w: int, region_size: Tuple[int]):
    bs, nhead, nregion, reg_size_square, head_dim = x.size()
    x = x.view(bs, nhead, region_h, region_w, region_size[0], region_size[1], head_dim)
    x = torch.einsum('bmhwpqd->bmdhpwq', x).reshape(bs, nhead * head_dim,
                                                    region_h * region_size[0], region_w * region_size[1])
    return x


def regional_routing_attention_torch(
        query: Tensor, key: Tensor, value: Tensor, scale: float,
        region_graph: LongTensor, region_size: Tuple[int],
        kv_region_size: Optional[Tuple[int]] = None,
        auto_pad=True) -> Tensor:
    kv_region_size = kv_region_size or region_size
    bs, nhead, q_nregion, topk = region_graph.size()

    q_pad_b, q_pad_r, kv_pad_b, kv_pad_r = 0, 0, 0, 0
    if auto_pad:
        _, _, Hq, Wq = query.size()
        q_pad_b = (region_size[0] - Hq % region_size[0]) % region_size[0]
        q_pad_r = (region_size[1] - Wq % region_size[1]) % region_size[1]
        if (q_pad_b > 0 or q_pad_r > 0):
            query = F.pad(query, (0, q_pad_r, 0, q_pad_b))

        _, _, Hk, Wk = key.size()
        kv_pad_b = (kv_region_size[0] - Hk % kv_region_size[0]) % kv_region_size[0]
        kv_pad_r = (kv_region_size[1] - Wk % kv_region_size[1]) % kv_region_size[1]
        if (kv_pad_r > 0 or kv_pad_b > 0):
            key = F.pad(key, (0, kv_pad_r, 0, kv_pad_b))
            value = F.pad(value, (0, kv_pad_r, 0, kv_pad_b))

    query, q_region_h, q_region_w = _grid2seq(query, region_size=region_size, num_heads=nhead)
    key, _, _ = _grid2seq(key, region_size=kv_region_size, num_heads=nhead)
    value, _, _ = _grid2seq(value, region_size=kv_region_size, num_heads=nhead)

    bs, nhead, kv_nregion, kv_region_size, head_dim = key.size()
    broadcasted_region_graph = region_graph.view(bs, nhead, q_nregion, topk, 1, 1). \
        expand(-1, -1, -1, -1, kv_region_size, head_dim)
    key_g = torch.gather(key.view(bs, nhead, 1, kv_nregion, kv_region_size, head_dim). \
                         expand(-1, -1, query.size(2), -1, -1, -1), dim=3,
                         index=broadcasted_region_graph)  # (bs, nhead, q_nregion, topk, kv_region_size, head_dim)
    value_g = torch.gather(value.view(bs, nhead, 1, kv_nregion, kv_region_size, head_dim). \
                           expand(-1, -1, query.size(2), -1, -1, -1), dim=3,
                           index=broadcasted_region_graph)  # (bs, nhead, q_nregion, topk, kv_region_size, head_dim)

    attn = (query * scale) @ key_g.flatten(-3, -2).transpose(-1, -2)
    attn = torch.softmax(attn, dim=-1)

    output = attn @ value_g.flatten(-3, -2)

    output = _seq2grid(output, region_h=q_region_h, region_w=q_region_w, region_size=region_size)

    if auto_pad and (q_pad_b > 0 or q_pad_r > 0):
        output = output[:, :, :Hq, :Wq]

    return output, attn


class BiLevelRoutingAttention_nchw(nn.Module):

    def __init__(self, dim, num_heads=8, n_win=7, qk_scale=None, topk=4, side_dwconv=3, auto_pad=False,
                 attn_backend='torch'):
        super().__init__()

        self.dim = dim
        self.num_heads = num_heads
        assert self.dim % num_heads == 0, 'dim must be divisible by num_heads!'
        self.head_dim = self.dim // self.num_heads
        self.scale = qk_scale or self.dim ** -0.5

        # 侧面深度卷积，用于局部上下文增强
        self.lepe = nn.Conv2d(dim, dim, kernel_size=side_dwconv, stride=1, padding=side_dwconv // 2,
                              groups=dim) if side_dwconv > 0 else \
            lambda x: torch.zeros_like(x)

        self.topk = topk
        self.n_win = n_win

        # 线性层用于生成查询（q）、键（k）和值（v）
        self.qkv_linear = nn.Conv2d(self.dim, 3 * self.dim, kernel_size=1)
        self.output_linear = nn.Conv2d(self.dim, self.dim, kernel_size=1)
        # 选择后端实现注意力机制
        if attn_backend == 'torch':
            self.attn_fn = regional_routing_attention_torch
        else:
            raise ValueError('CUDA implementation is not available yet. Please stay tuned.')

    def forward(self, x: Tensor, ret_attn_mask=False):

        N, C, H, W = x.size()
        region_size = (H // self.n_win, W // self.n_win)

        # 第一步：线性投影到q、k、v空间
        qkv = self.qkv_linear.forward(x)  # ncHW
        q, k, v = qkv.chunk(3, dim=1)  # ncHW

        # 第二步：区域到区域路由
        q_r = F.avg_pool2d(q.detach(), kernel_size=region_size, ceil_mode=True, count_include_pad=False)
        k_r = F.avg_pool2d(k.detach(), kernel_size=region_size, ceil_mode=True, count_include_pad=False)  # nchw
        q_r: Tensor = q_r.permute(0, 2, 3, 1).flatten(1, 2)  # n(hw)c
        k_r: Tensor = k_r.flatten(2, 3)  # nc(hw)
        a_r = q_r @ k_r  # n(hw)(hw), adj matrix of regional graph
        _, idx_r = torch.topk(a_r, k=self.topk, dim=-1)  # n(hw)k long tensor
        idx_r: LongTensor = idx_r.unsqueeze_(1).expand(-1, self.num_heads, -1, -1)

        # 第三步：非参数化的token-to-token注意力
        output, attn_mat = self.attn_fn(query=q, key=k, value=v, scale=self.scale,
                                        region_graph=idx_r, region_size=region_size
                                        )

        output = output + self.lepe(v)  # ncHW
        output = self.output_linear(output)  # ncHW

        if ret_attn_mask:
            return output, attn_mat

        return output
class MLP(nn.Module):
    def __init__(self, dim, embed_dim):
        super().__init__()
        self.proj = nn.Linear(dim, embed_dim)

    def forward(self, x: Tensor) -> Tensor:
        x = x.flatten(2).transpose(1, 2)
        x = self.proj(x)
        return x
class ConvModule(nn.Module):
    def __init__(self, c1, c2):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, 1, bias=False)
        self.bn = nn.BatchNorm2d(c2)        # use SyncBN in original
        self.activate = nn.ReLU(True)

    def forward(self, x: Tensor) -> Tensor:
        return self.activate(self.bn(self.conv(x)))
class SegFormerHead(nn.Module):
    def __init__(self, dims: list, embed_dim: int = 256, num_classes: int = 19):
        super().__init__()
        for i, dim in enumerate(dims):
            self.add_module(f"linear_c{i+1}", MLP(dim, embed_dim))

        self.linear_fuse = ConvModule(embed_dim*3, embed_dim)
        self.linear_pred = nn.Conv2d(embed_dim, num_classes, 1)
        self.dropout = nn.Dropout2d(0.1)

    def forward(self, features: Tuple[Tensor, Tensor, Tensor, Tensor]) -> Tensor:
        B, _, H, W = features[0].shape
        outs = [self.linear_c1(features[0]).permute(0, 2, 1).reshape(B, -1, *features[0].shape[-2:])]

        for i, feature in enumerate(features[1:]):
            cf = eval(f"self.linear_c{i+2}")(feature).permute(0, 2, 1).reshape(B, -1, *feature.shape[-2:])
            outs.append(F.interpolate(cf, size=(H, W), mode='bilinear', align_corners=False))

        seg = self.linear_fuse(torch.cat(outs[::-1], dim=1))
        seg = self.linear_pred(self.dropout(seg))
        return seg

import torch
import torch.nn as nn

# 定义全局上下文块类
class GlobalContextBlock(nn.Module):
    def __init__(self, inplanes, ratio, pooling_type="att", fusion_types=('channel_mul')) -> None:
        super().__init__()
        # 定义有效的融合类型
        valid_fusion_types = ['channel_add', 'channel_mul']
        # 断言池化类型为'avg'或'att'
        assert pooling_type in ['avg', 'att']
        # 断言至少使用一种融合方式
        assert len(fusion_types) > 0, 'at least one fusion should be used'
        # 初始化基本参数
        self.inplanes = inplanes
        self.ratio = ratio
        self.planes = int(inplanes * ratio)
        self.pooling_type = pooling_type
        self.fusion_type = fusion_types

        if pooling_type == 'att':
            self.conv_mask = nn.Conv2d(inplanes, 1, kernel_size=1)
            self.softmax = nn.Softmax(dim=2)
        else:
            # 否则，使用自适应平均池化
            self.avg_pool = nn.AdaptiveAvgPool2d(1)
        # 如果池化类型为'att'，使用1x1卷积作为掩码，并使用Softmax进行归一化
        if 'channel_add' in fusion_types:
            self.channel_add_conv = nn.Sequential(
                nn.Conv2d(self.inplanes, self.planes, kernel_size=1),
                nn.LayerNorm([self.planes, 1, 1]),
                nn.ReLU(inplace=True),
                nn.Conv2d(self.planes, self.inplanes, kernel_size=1)
            )
        else:
            self.channel_add_conv = None
        # 如果融合类型包含'channel_mul'，定义通道相乘卷积
        if 'channel_mul' in fusion_types:
            self.channel_mul_conv = nn.Sequential(
                nn.Conv2d(self.inplanes, self.planes, kernel_size=1),
                nn.LayerNorm([self.planes, 1, 1]),
                nn.ReLU(inplace=True),
                nn.Conv2d(self.planes, self.inplanes, kernel_size=1)
            )
        else:
            self.channel_mul_conv = None
        # 定义空间池化函数
    def spatial_pool(self, x):
        batch, channel, height, width = x.size()
        if self.pooling_type == 'att':
            input_x = x
            input_x = input_x.view(batch, channel, height * width) # 使用1x1卷积生成掩码
            input_x = input_x.unsqueeze(1)
            context_mask = self.conv_mask(x) # 使用1x1卷积生成掩码
            context_mask = context_mask.view(batch, 1, height * width)
            context_mask = self.softmax(context_mask)# 应用Softmax进行归一化
            context_mask = context_mask.unsqueeze(-1)
            context = torch.matmul(input_x, context_mask) # 计算上下文
            context = context.view(batch, channel, 1, 1)
        else:
            context = self.avg_pool(x) # 执行自适应平均池化
        return context

    # 定义前向传播函数
    def forward(self, x):
        context = self.spatial_pool(x)
        out = x
        if self.channel_mul_conv is not None:
            channel_mul_term = torch.sigmoid(self.channel_mul_conv(context))  # 将权重进行放大缩小
            out = out * channel_mul_term  # 与x进行相乘
        if self.channel_add_conv is not None:
            channel_add_term = self.channel_add_conv(context)
            out = out + channel_add_term
        return out



class CDO(torch.nn.Module):
    def __init__(self, device):
        super(CDO,self).__init__()
        self.out_size_h = 512
        self.out_size_w = 512
        self.device = device
        self.gamma = 2.
        self.OOM = True

        self.model = self.get_model()
        # self.modal_extractor = noisePrint_extractor()
        #
        # self.net_32 = GlobalContextBlock(32, ratio=1 / 16)
        # self.net_64 = GlobalContextBlock(64, ratio=1 / 16)
        # self.net_128 = GlobalContextBlock(128, ratio=1 / 16)


    def get_model(self) -> torch.nn.Module:
        backbone = "hrnet32"

        if backbone in valid_resnet_backbones:
            # expert network
            model_expert, _ = eval(f'{backbone}(pretrained=True)')
            # apprentice network
            model_apprentice, _ = eval(f'{backbone}(pretrained=False)')
        elif backbone in valid_hrnet_backbones:
            # expert network
            model_expert = HRNet_(backbone, pretrained=True)
            # apprentice network
            model_apprentice = HRNet_(backbone, pretrained=False)
        else:
            raise NotImplementedError

        # fix the parameters of the expert network
        # for param in model_expert.parameters():
        #     param.requires_grad = False

        # model_expert.eval()
        model = torch.nn.ModuleDict({'ME': model_expert, 'MA': model_apprentice})

        return model


    def forward(self, img, ori_img) -> dict:

        features = dict()
        # map the input to the expert domain and extract corresponding features
        with torch.no_grad():
            features['FE'] = self.model['ME'](img)

        # map the input to the apprentice domain and extract corresponding features
        features['FA'] = self.model['MA'](img)

        return features

    def save(self, path, metric):
        torch.save(self.model['MA'].state_dict(), path)

    def load(self, path):
        self.model['MA'].load_state_dict(torch.load(path, map_location=self.device))

    def train_mode(self):
        self.model['ME'].eval()
        self.model['MA'].train()

    def eval_mode(self):
        self.model['ME'].eval()
        self.model['MA'].eval()

    def cal_discrepancy(self, fe, fa, OOM, normal, gamma, aggregation=True):
        # normalize the features into uint vector
        fe = F.normalize(fe, p=2, dim=1)
        fa = F.normalize(fa, p=2, dim=1)

        # calculate feature-to-feature discrepancy d_p
        d_p = torch.sum((fe - fa) ** 2, dim=1)

        if OOM:
            # if OOM is utilized, we need to calculate the adaptive weights for individual features

            # calculate the mean discrepancy \mu_p to indicate the importance of individual features
            mu_p = torch.mean(d_p)

            if normal:
                # for normal samples: w = ((d_p) / \mu_p)^{\gamma}
                w = (d_p / mu_p) ** gamma

            else:
                # for abnormal samples: w = ((d_p) / \mu_p)^{-\gamma}
                w = (mu_p / d_p) ** gamma

            w = w.detach()

        else:
            # else, we manually assign each feature the same weight, i.e., 1
            w = torch.ones_like(d_p)

        if aggregation:
            d_p = torch.sum(d_p * w)

        sum_w = torch.sum(w)

        return d_p, sum_w


    def cal_loss(self, fe_list, fa_list, gamma=2, mask=None):
        loss = 0

        # interpolate the feature map into the size of the first hierarchy
        B, _, H_0, W_0 = fe_list[0].shape
        c = [32, 64, 128]
        for i in range(len(fe_list)):
            # fe_list[i] = self.atten(fe_list[i])
            # fa_list[i] = self.atten(fa_list[i])
            # fa_list[i] = eval(f"self.net_{c[i]}")(fa_list[i])
            # fe_list[i] = eval(f"self.net_{c[i]}")(fe_list[i])
            fe_list[i] = F.interpolate(fe_list[i], size=(H_0, W_0), mode='bilinear', align_corners=True)
            fa_list[i] = F.interpolate(fa_list[i], size=(H_0, W_0), mode='bilinear', align_corners=True)

        for fe, fa in zip(fe_list, fa_list):

            B, C, H, W = fe.shape

            # if mask is inputted, then we collaboratively optimize the discrepancies for normal and abnormal samples
            if mask is not None:
                mask_vec = F.interpolate(mask, (H, W), mode='nearest')
            else:
                mask_vec = torch.zeros((B, C, H, W))

            # reshape the mask, fe, and fa the the same shape for easily index
            mask_vec = mask_vec.permute(0, 2, 3, 1).reshape(-1, )

            fe = fe.permute(0, 2, 3, 1).reshape(-1, C)
            fa = fa.permute(0, 2, 3, 1).reshape(-1, C)

            # process normal and abnormal samples individually

            # normal features
            fe_n = fe[mask_vec == 0]
            fa_n = fa[mask_vec == 0]

            # synthetic abnormal features
            fe_s = fe[mask_vec != 0]
            fa_s = fa[mask_vec != 0]

            loss_n, weight_n = self.cal_discrepancy(fe_n, fa_n, OOM=self.OOM, normal=True, gamma=gamma,
                                                    aggregation=True)
            loss_s, weight_s = self.cal_discrepancy(fe_s, fa_s, OOM=self.OOM, normal=False, gamma=gamma,
                                                    aggregation=True)

            # L= {{\textstyle \sum_{i=1}^{N_n}}{{(w_n)_i}d(p_n)_i}-{\textstyle \sum_{j=1}^{N_a}}{{(w_s)_j}d(p_s)_j}} /
            # {{\textstyle \sum_{i=1}^{N_n}}{{(w_n)_i}}+{\textstyle \sum_{j=1}^{N_a}}{{(w_s)_j}}}
            loss += ((loss_n - loss_s) / (weight_n + weight_s) * B)

        return loss

    @torch.no_grad()
    def cal_am(self,shape ,**outputs):
        fe_list = outputs['FE']
        fa_list = outputs['FA']
        anomaly_map = 0
        c = [32,64,128]
        # interpolate the feature map into the size of the first hierarchy
        H_0, W_0 = fe_list[0].shape[2], fe_list[0].shape[3]
        for i in range(len(fe_list)):
            # fa_list[i] = eval(f"self.atten{c[i]}")(fa_list[i])
            # fe_list[i] = eval(f"self.atten{c[i]}")(fe_list[i])
            # fe_list[i] = self.atten(fe_list[i])
            # fa_list[i] = self.atten(fa_list[i])
            fe_list[i] = F.interpolate(fe_list[i], size=(H_0, W_0), mode='bilinear', align_corners=True)
            fa_list[i] = F.interpolate(fa_list[i], size=(H_0, W_0), mode='bilinear', align_corners=True)

        for fe, fa in zip(fe_list, fa_list):
            _, _, h, w = fe.shape
            a_map, _ = self.cal_discrepancy(fe, fa, gamma=self.gamma, aggregation=False, OOM=True,
                                            normal=True)

            a_map = torch.unsqueeze(a_map, dim=1)
            a_map = F.interpolate(a_map, size=shape, mode='bilinear', align_corners=False)

            # fuse anomaly maps from different hierarchies
            anomaly_map += a_map
        # anomaly_map = self.atten(anomaly_map)
        am_np = anomaly_map.squeeze(1).cpu().numpy()
        # am_np = anomaly_map.squeeze(1)

        am_np_list = []

        for i in range(am_np.shape[0]):
            am_np[i] = gaussian_filter(am_np[i], sigma=4)
            am_np_list.append(am_np[i])

        return am_np_list






